﻿using System.Collections.Generic;
using System.Linq;
using Shell.Common;

namespace Shell
{
    public abstract class Node
    {
    }

    public abstract class BranchNode : Node
    {
        public struct WeightedNode
        {
            public Node Node;
            public double Weight;
        }

        public abstract IEnumerable<WeightedNode> Process();

        protected static IEnumerable<WeightedNode> Identity(IEnumerable<Node> nodes)
        {
            return nodes.Select(_ => new WeightedNode
                {
                    Node = _,
                    Weight = 1
                });
        }
    }

    public abstract class LeafNode : Node
    {
    }

    
    public class Searcher
    {
        public class NodeContext
        {
            public NodeContext Parent;
            public BranchNode Node;
            public double Weight;
        }

        public IEnumerable<TResult> Search<TResult>(BranchNode initialGoal)
            where TResult : LeafNode
        {
            return Search<TResult>(initialGoal.Yield());
        }

        public IEnumerable<TResult> Search<TResult>(IEnumerable<BranchNode> initialNodes)
            where TResult : LeafNode
        {
            var initialContexts = initialNodes.Select(node => new NodeContext
                {
                    Parent =  null,
                    Node = node,
                    Weight = 1
                });

            var contexts = new Stack<NodeContext>(initialContexts);

            while (contexts.Count > 0)
            {
                var current = contexts.Pop();
                var result = current.Node.Process().ToList();
                var nextNodes = ExtractBranches(result, current);
                var resultNodes = ExtractLeafs<TResult>(result);

                contexts.PushRange(nextNodes);
                foreach (var r in resultNodes) yield return r;
            }
        }

        private static IEnumerable<TResult> ExtractLeafs<TResult>(IEnumerable<BranchNode.WeightedNode> result)
            where TResult : LeafNode
        {
            return result
                .Where(_ => _.Node is TResult)
                .Select(_ => _.Node as TResult);
        }

        private static IEnumerable<NodeContext> ExtractBranches(IEnumerable<BranchNode.WeightedNode> result, NodeContext current)
        {
            var nextNodes = result
                .Where(_ => _.Node is BranchNode)
                .Select(_ => new NodeContext
                    {
                        Parent = current,
                        Node = _.Node as BranchNode,
                        Weight = current.Weight*_.Weight
                    });
            return nextNodes;
        }
    }
}